// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: daniel

#pragma once

#include <memory>
#include <optional>
#include <string>
#include <unordered_map>
#include <utility>

#include "coin2/exchange/base/order/rate_config.h"
#include "coin2/exchange/base/order/rate_tracker.h"

constexpr auto GET_METHOD_PER_IP = "get_method_per_ip";
constexpr auto POST_METHOD_PER_IP = "post_method_per_ip";
constexpr int64_t NS = 1'000'000'000LL;
constexpr auto EXAMPLE_FUTURES_BYBIT = R"(
"connection_config": {
  "key_filepath": "/home/ec2-user/config/trade/xunke06/bybit_futures/trade_key.json",
  "rate_limit_info": {
    "futures_bybit": {
      "rate_limits": {
        "get_method_per_ip": {
          "secs": 1,
          "reqs": 70
        },
        "post_method_per_ip": {
          "secs": 1,
          "reqs": 50
        },
        "group_query_account": {
          "secs": 60,
          "reqs": 120
        },
        "group_query_order": {
          "secs": 60,
          "reqs": 600
        },
        "group_query_trade": {
          "secs": 60,
          "reqs": 120
        },
        "group_trade_order": {
          "secs": 60,
          "reqs": 100
        }
      }
    }
  }
})";

class BybitRateTracker : public coin2::exchange::base::order::IRateTracker {
 public:
  using RateTracker = coin2::exchange::base::order::RateTracker;
  using ActionType = coin2::exchange::base::order::IRateTracker::ActionType;
  using HttpContextPtr = coin2::exchange::base::order::IRateTracker::HttpContextPtr;
  using OrderLoggerCB = coin2::exchange::base::order::IRateTracker::OrderLoggerCB;
  using RateLimitReportWriter = coin2::exchange::base::order::RateLimitReportWriter;
  using FixedWindowRateTracker = coin2::exchange::base::order::FixedWindowRateTracker;
  using RateLimitInfoItem = coin2::exchange::order::RateLimitInfoItem;
  using RateLimitInfoFuturesBybit = coin2::exchange::order::RateLimitInfoFuturesBybit;
  using OrderType = coin::proto::OrderType;

  struct RatioValveConfig {
    double ratio_valve_query_order = 1.0;
    double ratio_valve_maker_order = 1.0;
    double ratio_valve_taker_order = 1.0;
    double ratio_valve_amend_order = 1.0;
    double ratio_valve_cancel_order = 1.0;
    RatioValveConfig() = default;

    static auto ExtractRatioValveConfig(const coin2::exchange::order::RateLimitInfoItem& config)
        -> RatioValveConfig {
      RatioValveConfig ret;
      for (auto& [key, val] : config.ratio_valve_config()) {
        if (key == "query_order") {
          ret.ratio_valve_query_order = val;
        } else if (key == "maker_order") {
          ret.ratio_valve_maker_order = val;
        } else if (key == "taker_order") {
          ret.ratio_valve_taker_order = val;
        } else if (key == "amend_order") {
          ret.ratio_valve_amend_order = val;
        } else if (key == "cancel_order") {
          ret.ratio_valve_cancel_order = val;
        }
      }
      return ret;
    }

    friend std::ostream& operator<<(std::ostream& out, const RatioValveConfig& ratio_valve_config) {
      out << "\nquery_order_ratio_valve=" << ratio_valve_config.ratio_valve_query_order
          << "\nmaker_order_ratio_valve=" << ratio_valve_config.ratio_valve_maker_order
          << "\ntaker_order_ratio_valve=" << ratio_valve_config.ratio_valve_taker_order
          << "\namend_order_ratio_valve=" << ratio_valve_config.ratio_valve_amend_order
          << "\ncancel_order_ratio_valve=" << ratio_valve_config.ratio_valve_cancel_order;
      return out;
    }
  };

  class BybitRateTrackerImpl : public coin2::exchange::base::order::FixedWindowRateTracker {
   public:
    BybitRateTrackerImpl(
        int64_t window_size,
        const RatioValveConfig ratio_valve_config,
        const std::string& name)
        : coin2::exchange::base::order::FixedWindowRateTracker{window_size},
          ratio_valve_config_{ratio_valve_config},
          name_{name} {
      DLOG(INFO) << ratio_valve_config_;
    }

    bool HasEnoughRateLimit(int64_t timestamp, ActionType action_type, bool is_maker_order) {
      auto max = MaxReq();
      auto rem = GetRemaining(timestamp);
      auto sum = max - rem;
      auto weight = ratio_valve_config_.ratio_valve_query_order;
      switch (action_type) {
        case ActionType::QUERY:
          weight = ratio_valve_config_.ratio_valve_query_order;
          break;
        case ActionType::PLACE:
          weight = is_maker_order ? ratio_valve_config_.ratio_valve_maker_order
                                  : ratio_valve_config_.ratio_valve_taker_order;
          break;
        case ActionType::AMEND:
          weight = ratio_valve_config_.ratio_valve_amend_order;
          break;
        case ActionType::CANCEL:
          weight = ratio_valve_config_.ratio_valve_cancel_order;
          break;
        default:
          // shouldn't reach here
          break;
      }
      double used_weight = static_cast<double>(sum) / max;

      ///////////
      // std::string action_type_str = "";
      // switch (action_type) {
      //   case ActionType::QUERY:
      //     action_type_str = "QUERY";
      //     break;
      //   case ActionType::PLACE:
      //     action_type_str = "PLACE";
      //     break;
      //   case ActionType::AMEND:
      //     action_type_str = "AMEND";
      //     break;
      //   case ActionType::CANCEL:
      //     action_type_str = "CANCEL";
      //     break;
      //   default:
      //     break;
      // }
      // std::string order_type_str = "";
      // if (is_maker_order) {
      //   order_type_str = "order_maker";
      // } else {
      //   order_type_str = "order_taker";
      // }
      // DLOG(INFO) << "name: " << name_ << ", info: num_rate_limit=" << max
      //            << ", num_remain=" << rem << ", num_used=" << sum
      //            << ", used_ratio=" << used_weight << ", limit_ratio=" << weight
      //            << ", action_type=" << action_type_str << ", order_type=" << order_type_str;
      ///////////

      return used_weight < weight;
    }

   private:
    RatioValveConfig ratio_valve_config_;
    const std::string name_;
  };

  class IPRateTracker {
   public:
    explicit IPRateTracker(const RateLimitInfoFuturesBybit& rate_limit_config) {
      for (auto& [key, val] : rate_limit_config.rate_limits()) {
        if (key == GET_METHOD_PER_IP) {
          get_method_rate_tracker_ = std::make_unique<BybitRateTrackerImpl>(
              val.secs() * NS,
              RatioValveConfig::ExtractRatioValveConfig(val),
              key);
          get_method_rate_tracker_->SetRateLimitInfo(val.reqs(), val.reqs());
        } else if (key == POST_METHOD_PER_IP) {
          post_method_rate_tracker_ = std::make_unique<BybitRateTrackerImpl>(
              val.secs() * NS,
              RatioValveConfig::ExtractRatioValveConfig(val),
              key);
          post_method_rate_tracker_->SetRateLimitInfo(val.reqs(), val.reqs());
        }
      }
      Validate();
    }

    void Validate() const {
      CHECK(get_method_rate_tracker_) << "missing get method rate tracker for Bybit";
      CHECK(post_method_rate_tracker_) << "missing post method rate tracker for Bybit";
    }

    void Push(int32_t num_reqs, int64_t timestamp, ActionType action_type) {
      if (action_type == ActionType::QUERY) {
        get_method_rate_tracker_->Push(num_reqs, timestamp);
      } else {
        post_method_rate_tracker_->Push(num_reqs, timestamp);
      }
    }

    bool HasEnoughRateLimit(int64_t ts, ActionType action_type, bool is_maker_order) {
      bool ret = true;
      if (action_type == ActionType::QUERY) {
        ret = get_method_rate_tracker_->HasEnoughRateLimit(ts, action_type, is_maker_order);
      } else {
        ret = post_method_rate_tracker_->HasEnoughRateLimit(ts, action_type, is_maker_order);
      }
      return ret;
    }

   private:
    std::unique_ptr<BybitRateTrackerImpl> get_method_rate_tracker_ = {};
    std::unique_ptr<BybitRateTrackerImpl> post_method_rate_tracker_ = {};
  };

  class SymbolRateTracker {
   public:
    explicit SymbolRateTracker(const RateLimitInfoFuturesBybit& rate_limit_config) {
      for (auto& [key, val] : rate_limit_config.rate_limits()) {
        if (key.empty() || key == GET_METHOD_PER_IP || key == POST_METHOD_PER_IP) {
          continue;
        }
        rate_tracker_map_.emplace(
            key,
            std::make_unique<BybitRateTrackerImpl>(
                val.secs() * NS,
                RatioValveConfig::ExtractRatioValveConfig(val),
                key));
        rate_tracker_map_[key]->SetRateLimitInfo(val.reqs(), val.reqs());
      }
      Validate();
    }

    void Validate() { /* TODO(daniel): validate */
    }

    bool HasEnoughRateLimit(
        int64_t ts,
        ActionType action_type,
        bool is_maker_order,
        const std::string& endpoint) {
      if (rate_tracker_map_.count(endpoint) == 0) {
        return true;
      }
      auto&& rate_tracker = rate_tracker_map_.at(endpoint);
      return rate_tracker->HasEnoughRateLimit(ts, action_type, is_maker_order);
    }

    void Push(int32_t num_reqs, int64_t ts, ActionType action_type, const std::string& endpoint) {
      if (rate_tracker_map_.count(endpoint) == 0) {
        return;
      }
      auto&& rate_tracker = rate_tracker_map_.at(endpoint);
      rate_tracker->Push(num_reqs, ts);
    }

   private:
    std::unordered_map<std::string, std::unique_ptr<BybitRateTrackerImpl>> rate_tracker_map_;
  };

 public:
  explicit BybitRateTracker(
      const BybitFuturesNativePrivateClient& rest_client,
      const MarketExchangeApi& mea,
      const coin2::exchange::order::OrderConnectionConfig& cfg)
      : rest_client_{rest_client} {
    using RateConfig = coin2::exchange::base::order::RateConfig;
    RateConfig::Get()->GetConfig(mea, cfg, &rate_limit_info_);
    // validate config
    Validate(mea);
    // ip rate tracker
    ip_rate_tracker_ = std::make_unique<IPRateTracker>(rate_limit_info_);
    // rate limit report
    rate_limit_report_writer_ = std::make_unique<RateLimitReportWriter>();
    rate_limit_report_writer_->SetRateLimitInfoString(&rate_limit_info_);
  }

  void Validate(const MarketExchangeApi& mea) const {
    // check mandatory field
    if (rate_limit_info_.rate_limits().count(GET_METHOD_PER_IP) == 0 ||
        rate_limit_info_.rate_limits().count(POST_METHOD_PER_IP) == 0 ||
        rate_limit_info_.rate_limits().count(rest_client_.GetEndpointGroupQueryAccount()) == 0 ||
        rate_limit_info_.rate_limits().count(rest_client_.GetEndpointGroupQueryOrder()) == 0 ||
        rate_limit_info_.rate_limits().count(rest_client_.GetEndpointGroupQueryTrade()) == 0 ||
        rate_limit_info_.rate_limits().count(rest_client_.GetEndpointGroupTradeOrder()) == 0) {
      using RateConfig = coin2::exchange::base::order::RateConfig;
      RateConfig::CheckFail(mea.String(), EXAMPLE_FUTURES_BYBIT);
      CHECK(false) << "invalid rate limit config for Bybit";
    }
    // print rate limit info
    std::stringstream ss;
    ss << "\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~"
       << "\nOG request rate was set as below:\n";
    for (auto [key, val] : rate_limit_info_.rate_limits()) {
      ss << "\n" << key << ": [" << val.reqs() << " (reqs) / " << val.secs() << " (seconds)]";
    }
    ss << "\n\nOr you can configure customized rate in driver config. For example: "
       << EXAMPLE_FUTURES_BYBIT
       << "\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~";
    LOG(INFO) << ss.str();
  }

  void SetRateLimitInfo(int32_t rate_limit, int32_t rate_limit_remaining) override { CHECK(false); }
  void Push(int32_t num_reqs, int64_t timestamp, ActionType action_type) override { CHECK(false); }
  bool HasEnoughRateLimit(int64_t timestamp, ActionType action_type) override { CHECK(false); }
  void ParseRateLimitInfo(const HttpContextPtr&, ActionType) override { /* do nothing */
  }

  void SetRateLimitInfoString(const google::protobuf::Message* message) override {
    rate_limit_report_writer_->SetRateLimitInfoString(message);
  }

  void WriteRateLimitReportMessage(OrderLoggerCB callback) override {
    if (!callback) {
      return;
    }
    rate_limit_report_writer_->Write(callback);
  }

  double GetMaxRateUsedRatio(int64_t timestamp, ActionType action_type) override {
    // TODO(daniel):
    return 0.;
  }

  bool HasEnoughRateLimit(
      int64_t ts,
      ActionType action_type,
      const std::string& symbol,
      const std::string& endpoint) {
    bool ret = true;
    // check symbol rate limit
    if (rate_tracker_map_.count(symbol) == 0) {
      rate_tracker_map_.emplace(symbol, std::make_unique<SymbolRateTracker>(rate_limit_info_));
      ret = true;
    } else {
      auto& rate_tracker = rate_tracker_map_.at(symbol);
      ret = rate_tracker->HasEnoughRateLimit(ts, action_type, is_maker_order_, endpoint);
    }
    // check ip rate limit
    if (ret) {
      ret = ip_rate_tracker_->HasEnoughRateLimit(ts, action_type, is_maker_order_);
    }
    if (!ret) {
      rate_limit_report_writer_->SetHitRateLimit(ts, action_type);
    }

    return ret;
  }

  void Push(
      int32_t num_reqs,
      int64_t ts,
      ActionType action_type,
      const std::string& symbol,
      const std::string& endpoint) {
    // update ip rate limit
    ip_rate_tracker_->Push(num_reqs, ts, action_type);
    // update symbol rate limit
    if (rate_tracker_map_.count(symbol) == 0) {
      return;
    }
    auto&& rate_tracker = rate_tracker_map_.at(symbol);
    rate_tracker->Push(num_reqs, ts, action_type, endpoint);
  }

  void SetMakerOrder(bool is_maker_order) { is_maker_order_ = is_maker_order; }

 private:
  const BybitFuturesNativePrivateClient& rest_client_;
  coin2::exchange::order::RateLimitInfoFuturesBybit rate_limit_info_;
  std::unique_ptr<IPRateTracker> ip_rate_tracker_ = {};
  std::unordered_map<std::string, std::unique_ptr<SymbolRateTracker>> rate_tracker_map_;
  std::unique_ptr<RateLimitReportWriter> rate_limit_report_writer_ = {};
  bool is_maker_order_ = false;
};
